#include <efi.h>

#define	MAP_SIZE			0x40000000
#define	MINI_MEMBLOCK			8

EFI_SYSTEM_TABLE			*st;
EFI_BOOT_SERVICES			*bs;
EFI_STATUS				success;

void print(CHAR16 str[])
{
	st->ConOut->OutputString(st->ConOut, str);
}

void checkSuccess(CHAR16 str[])
{
	if(success != EFI_SUCCESS)
	{
		print(L"mmap error: ");
		print(str);
		print(L"\r\n");
		for(;;){}
	}
}

#include "printhex.c"

void displayMapEntry(EFI_MEMORY_DESCRIPTOR *descPtr)
{

	printH32(descPtr->Type);
	print(L" ");
	printH64(descPtr->PhysicalStart);
	print(L" ");
	printH64(descPtr->VirtualStart);
	print(L" ");
	printH64(descPtr->NumberOfPages);
	print(L" ");
	printH64(descPtr->Attribute);
	print(L"\r\n");
}

EFI_STATUS efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *ST)
{
	void					*memBlockAddress;
	long long unsigned int			size, descSize, key, mapLoop;
	unsigned int 				descVersion;
	EFI_MEMORY_DESCRIPTOR			*mapPtr, *descPtr;
	uint8_t					miniMem[MINI_MEMBLOCK];
	EFI_INPUT_KEY				iKey;

	/* save in global vars */
	st = ST;
	bs = ST->BootServices;

	/* Get mem map just to get the required size! */
	size = MINI_MEMBLOCK;
	bs->GetMemoryMap(&size, &miniMem, &key, &descSize, &descVersion);
	/* No matter if the function worked, now the var 'size' contains
		the required size for the mem map */
	
	/* alloc mem for the mem map */
	success = bs->AllocatePool(0, size, &memBlockAddress);
	checkSuccess(L"Can't allocate");
	
	/* get mem map */
	mapPtr = memBlockAddress; /* type casting from void*
				to EFI_MEMORY_DESCRIPTOR* */
	success = bs->GetMemoryMap(&size, mapPtr, &key, &descSize,
		&descVersion);
	checkSuccess(L"No mem map");
	
	/* Go through mem map */
	print(L"Number of descriptors: ");
	printH64(size/descSize);
	print(L"\r\n");
	mapLoop = 0;
	do
	{
		descPtr = mapPtr + (mapLoop * descSize);
		if(descPtr->NumberOfPages != 0)
		{
			displayMapEntry(descPtr);
			do
			{
			}
			while(bs->CheckEvent(st->ConIn->WaitForKey) != EFI_SUCCESS);
			mapLoop++;
		}
	}
	while(mapLoop<=size/descSize && descPtr->NumberOfPages != 0);

	/* We're finished */	
	for(;;){}
	return EFI_SUCCESS; /* Make compiler happy */
}

